# @Time : 2021/8/4 8:51
# @Author : Li Kunlun
# @Description : 线性回归的简洁实现

from mxnet import autograd, nd

# 1、生成数据集
num_inputs = 2
num_examples = 1000
true_w = [2, -3.4]
true_b = 4.2
features = nd.random.normal(scale=1, shape=(num_examples, num_inputs))
labels = true_w[0] * features[:, 0] + true_w[1] * features[:, 1] + true_b
labels += nd.random.normal(scale=0.01, shape=labels.shape)

from mxnet.gluon import data as gdata

# 2、读取数据集 Gluon提供了data包来读取数据
batch_size = 10
# 将训练数据的特征和标签组合
dataset = gdata.ArrayDataset(features, labels)
# 随机读取小批量
data_iter = gdata.DataLoader(dataset, batch_size, shuffle=True)

for X, y in data_iter:
    print(X, y)
    break

from mxnet.gluon import nn

# 3、定义模型
net = nn.Sequential()
net.add(nn.Dense(1))

from mxnet import init

# 4、初始化模型参数
net.initialize(init.Normal(sigma=0.01))

from mxnet.gluon import loss as gloss

# 5、定义损失函数
loss = gloss.L2Loss()  # 平方损失又称L2范数损失

from mxnet import gluon

# 6、定义优化算法
# 该优化算法将用来迭代net实例所有通过add函数嵌套的层所包含的全部参数。这些参数可以通过collect_params函数获取。
trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': 0.03})

# 训练模型
num_epochs = 3
for epoch in range(1, num_epochs + 1):
    for X, y in data_iter:
        with autograd.record():
            l = loss(net(X), y)
        l.backward()

        # l.shape->  (10,)
        # print("l.shape-> ",l.shape)

        trainer.step(batch_size)
    l = loss(net(features), labels)
    print('epoch %d, loss: %f' % (epoch, l.mean().asnumpy()))

# 输出结果和原始结果相差不大
dense = net[0]
print(true_w, dense.weight.data())
print(true_b, dense.bias.data())

print("true_w->", true_w, "w->", dense.weight.data())
print("true_b->", true_b, "b->", dense.bias.data())
